import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
import cv2

BasePath = 'D:/360 dataset/360_Saliency_dataset_2018ECCV/dataset/'
SavePath = 'D:/360 dataset/360_Saliency_dataset_2018ECCV/sample_data/'
npyPath = 'D:/360 dataset/360_Saliency_dataset_2018ECCV/'
indexList = [i for i in range(217, 321)]

info = pickle.load(open(BasePath+r'\vinfo.pkl', 'rb'))

print(type(info))

step = 30
shape = (512, 256)
saliency = []
start = 1

for idx in indexList:
    nums = 0
    fileList = os.listdir(BasePath+str(idx))
    for file in fileList:
        if os.path.splitext(file)[-1] == '.jpg':
            nums += 1

    for i in range(1, nums+1, step):
        if i < 100:
            continue
        if i<10:
            i = '000'+str(i)
        elif 10 <= i < 100:
            i ='00'+str(i)
        elif 100 <= i < 1000:
            i = '0'+str(i)
        else:
            i = str(i)

        image = cv2.imread(BasePath+str(idx)+'/'+i+'.jpg')
        cv2.imwrite(SavePath+str(start)+'.jpg', cv2.resize(image, shape))
        start += 1

        gt = np.load(BasePath+str(idx)+'/'+i+'_gt.npy')[0]
        saliency.append(gt)

    print(f'sample {idx} fileList finish with {start-1}.jpg', len(saliency))

np.save(npyPath+'sample_data.npy', np.array(saliency))
